// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/array/builder_primitive.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <utility>
#include <vector>

#include <nebula/core/array.h>
#include <nebula/core/buffer.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <nebula/types/type_traits.h>
#include <nebula/bits/bit_util.h>
#include <nebula/numeric/int_util.h>
#include <turbo/log/logging.h>

namespace nebula {

    // ----------------------------------------------------------------------
    // Null builder

    turbo::Status NullBuilder::finish_internal(std::shared_ptr<ArrayData> *out) {
        *out = ArrayData::create(null(), length_, {nullptr}, length_);
        length_ = null_count_ = 0;
        return turbo::OkStatus();
    }

    BooleanBuilder::BooleanBuilder(MemoryPool *pool, int64_t alignment)
            : ArrayBuilder(pool, alignment), data_builder_(pool, alignment) {}

    BooleanBuilder::BooleanBuilder(const std::shared_ptr<DataType> &type, MemoryPool *pool,
                                   int64_t alignment)
            : BooleanBuilder(pool, alignment) {
        KCHECK_EQ(Type::BOOL, type->id());
    }

    void BooleanBuilder::reset() {
        ArrayBuilder::reset();
        data_builder_.reset();
    }

    turbo::Status BooleanBuilder::resize(int64_t capacity) {
        TURBO_RETURN_NOT_OK(check_capacity(capacity));
        capacity = std::max(capacity, kMinBuilderCapacity);
        TURBO_RETURN_NOT_OK(data_builder_.resize(capacity));
        return ArrayBuilder::resize(capacity);
    }

    turbo::Status BooleanBuilder::finish_internal(std::shared_ptr<ArrayData> *out) {
        TURBO_MOVE_OR_RAISE(auto null_bitmap, null_bitmap_builder_.FinishWithLength(length_));
        TURBO_MOVE_OR_RAISE(auto data, data_builder_.FinishWithLength(length_));

        *out = ArrayData::create(boolean(), length_, {null_bitmap, data}, null_count_);

        capacity_ = length_ = null_count_ = 0;
        return turbo::OkStatus();
    }

    turbo::Status BooleanBuilder::append_values(const uint8_t *values, int64_t length,
                                               const uint8_t *valid_bytes) {
        TURBO_RETURN_NOT_OK(Reserve(length));

        int64_t i = 0;
        data_builder_.unsafe_append<false>(length,
                                          [values, &i]() -> bool { return values[i++] != 0; });
        ArrayBuilder::unsafe_append_to_bitmap(valid_bytes, length);
        return turbo::OkStatus();
    }

    turbo::Status BooleanBuilder::append_values(const uint8_t *values, int64_t length,
                                               const uint8_t *validity, int64_t offset) {
        TURBO_RETURN_NOT_OK(Reserve(length));
        data_builder_.unsafe_append(values, offset, length);
        ArrayBuilder::unsafe_append_to_bitmap(validity, offset, length);
        return turbo::OkStatus();
    }

    turbo::Status BooleanBuilder::append_values(const uint8_t *values, int64_t length,
                                               const std::vector<bool> &is_valid) {
        TURBO_RETURN_NOT_OK(Reserve(length));
                DKCHECK_EQ(length, static_cast<int64_t>(is_valid.size()));
        int64_t i = 0;
        data_builder_.unsafe_append<false>(length,
                                          [values, &i]() -> bool { return values[i++]; });
        ArrayBuilder::unsafe_append_to_bitmap(is_valid);
        return turbo::OkStatus();
    }

    turbo::Status BooleanBuilder::append_values(const std::vector<uint8_t> &values,
                                               const std::vector<bool> &is_valid) {
        return append_values(values.data(), static_cast<int64_t>(values.size()), is_valid);
    }

    turbo::Status BooleanBuilder::append_values(const std::vector<uint8_t> &values) {
        return append_values(values.data(), static_cast<int64_t>(values.size()));
    }

    turbo::Status BooleanBuilder::append_values(const std::vector<bool> &values,
                                               const std::vector<bool> &is_valid) {
        const int64_t length = static_cast<int64_t>(values.size());
        TURBO_RETURN_NOT_OK(Reserve(length));
                DKCHECK_EQ(length, static_cast<int64_t>(is_valid.size()));
        int64_t i = 0;
        data_builder_.unsafe_append<false>(length,
                                          [&values, &i]() -> bool { return values[i++]; });
        ArrayBuilder::unsafe_append_to_bitmap(is_valid);
        return turbo::OkStatus();
    }

    turbo::Status BooleanBuilder::append_values(const std::vector<bool> &values) {
        const int64_t length = static_cast<int64_t>(values.size());
        TURBO_RETURN_NOT_OK(Reserve(length));
        int64_t i = 0;
        data_builder_.unsafe_append<false>(length,
                                          [&values, &i]() -> bool { return values[i++]; });
        ArrayBuilder::unsafe_set_not_null(length);
        return turbo::OkStatus();
    }

    turbo::Status BooleanBuilder::append_values(int64_t length, bool value) {
        TURBO_RETURN_NOT_OK(Reserve(length));
        data_builder_.unsafe_append(length, value);
        ArrayBuilder::unsafe_set_not_null(length);
        return turbo::OkStatus();
    }

}  // namespace nebula
